import scipy.sparse
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from models.BaseModel import GeneralModel
import scipy


#item_rep = torch.eye(n_i).to(device) 
#user_rep = torch.zeros(n_u,n_i).to(device)
#for i in range(k):
#    print("Running layer", i)
#    user_rep_temp = torch.sparse.mm(adj,item_rep) + user_rep
#    item_rep_temp = torch.sparse.mm(adj.transpose(0,1),user_rep) + item_rep
#    user_rep = user_rep_temp
#    item_rep = item_rep_temp

#记录item和user的数量n_u,n_i
#构建n_u*n_i大小的H1（我们的模型）
#构建n_i*n_i大小的H2（每个item对应每个用户的初始评分）
#生成一个初始图A（通过训练数据train矩阵得到，相当于是描述了train中user和item的交互情况）
#反复通过H1(1) = A * H2(0) + H1(0), H2(1) = A * H1(0) + H2(0)
#得到H1(L) = A * H2(L-1) + H1(L-1), H2(L) = A * H1(L-1) + H2(L-1)
class GTE(GeneralModel):
    reader = 'BaseReader'  # assign a reader class, BaseReader by default
    runner = 'GTERunner'  # assign a runner class, BaseRunner by default

    @staticmethod
    def parse_model_args(parser):
        parser.add_argument('--k', type=int, default=3,
                            help='The number of iteration times.')
        return GeneralModel.parse_model_args(parser)
    
    def __init__(self, args, corpus):
        super().__init__(args, corpus)
        self.raw_data = corpus.data_df['train']
        one_list = [1.0]*len(self.raw_data['time'])
        self.data = scipy.sparse.coo_matrix((self.raw_data['time'],(self.raw_data['user_id'],self.raw_data['item_id'])), shape=(self.user_num,self.item_num))
        self._define_params()
        self.apply(self.init_weights)
        for i in range(self.k):
            print("Running layer", i)
            user_rep_temp = torch.sparse.mm(self.adj,self.item_h) + self.user_h
            item_rep_temp = torch.sparse.mm(self.adj.transpose(0,1),self.user_h) + self.item_h
            self.user_h = user_rep_temp
            self.item_h = item_rep_temp
        self.model = self.user_h.cpu().numpy()

    def get_predict(self, predictions):
        train_csr = (self.data!=0).astype(np.float32)
        mask = train_csr.toarray()
        preds = predictions * (1-mask)
        predictions = (-preds).argsort()
        return predictions

    def _define_params(self):
        # define parameters in the model
        self.item_h = torch.eye(self.item_num).to(self.device)
        self.user_h = torch.zeros(self.user_num, self.item_num).to(self.device)
        self.adj = self.scipy_sparse_mat_to_torch_sparse_tensor(self.data).coalesce().to(self.device)
        self.k = 3

    def scipy_sparse_mat_to_torch_sparse_tensor(self, sparse_mx):
        sparse_mx = sparse_mx.tocoo().astype(np.float32)
        indices = torch.from_numpy(
            np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
        values = torch.from_numpy(sparse_mx.data)
        shape = torch.Size(sparse_mx.shape)
        return torch.sparse.FloatTensor(indices, values, shape)

    def forward(self, feed_dict):
        pass

    class Dataset(GeneralModel.Dataset):
        # construct feed_dict for a single instance (called by __getitem__)
        # will be collated to a integrated feed dict for each batch
        def _get_feed_dict(self, index):
            feed_dict = super()._get_feed_dict(index)
            return feed_dict